package rate_limit

import (
	"sync"
	"time"

	"apiGateWay/internal/config"
	"apiGateWay/pkg/logger"

	"github.com/gin-gonic/gin"
)

// Limiter 限流器接口
type Limiter interface {
	Allow() bool
	Wait() time.Duration
}

// TokenBucket 令牌桶限流器
type TokenBucket struct {
	capacity   int
	tokens     int
	refillRate float64
	lastRefill time.Time
	mu         sync.Mutex
}

// NewTokenBucket 创建令牌桶
func NewTokenBucket(capacity, qps int) *TokenBucket {
	return &TokenBucket{
		capacity:   capacity,
		tokens:     capacity,
		refillRate: float64(qps),
		lastRefill: time.Now(),
	}
}

// Allow 检查是否允许请求
func (tb *TokenBucket) Allow() bool {
	tb.mu.Lock()
	defer tb.mu.Unlock()

	now := time.Now()
	elapsed := now.Sub(tb.lastRefill).Seconds()
	tokensToAdd := int(elapsed * tb.refillRate)

	if tokensToAdd > 0 {
		tb.tokens = min(tb.capacity, tb.tokens+tokensToAdd)
		tb.lastRefill = now
	}

	if tb.tokens > 0 {
		tb.tokens--
		return true
	}

	return false
}

// Wait 等待直到允许请求
func (tb *TokenBucket) Wait() time.Duration {
	if tb.Allow() {
		return 0
	}

	tb.mu.Lock()
	defer tb.mu.Unlock()

	// 计算需要等待的时间
	tokensNeeded := 1
	waitTime := time.Duration(float64(tokensNeeded)/tb.refillRate) * time.Second
	return waitTime
}

// LeakyBucket 漏桶限流器
type LeakyBucket struct {
	capacity   int
	leakRate   float64
	queue      []time.Time
	lastLeak   time.Time
	mu         sync.Mutex
}

// NewLeakyBucket 创建漏桶
func NewLeakyBucket(capacity, qps int) *LeakyBucket {
	return &LeakyBucket{
		capacity: capacity,
		leakRate: float64(qps),
		lastLeak: time.Now(),
	}
}

// Allow 检查是否允许请求
func (lb *LeakyBucket) Allow() bool {
	lb.mu.Lock()
	defer lb.mu.Unlock()

	now := time.Now()
	elapsed := now.Sub(lb.lastLeak).Seconds()
	leaks := int(elapsed * lb.leakRate)

	// 移除过期的请求
	if leaks > 0 && len(lb.queue) > 0 {
		if leaks >= len(lb.queue) {
			lb.queue = lb.queue[:0]
		} else {
			lb.queue = lb.queue[leaks:]
		}
		lb.lastLeak = now
	}

	if len(lb.queue) < lb.capacity {
		lb.queue = append(lb.queue, now)
		return true
	}

	return false
}

// Wait 等待直到允许请求
func (lb *LeakyBucket) Wait() time.Duration {
	if lb.Allow() {
		return 0
	}

	// 计算需要等待的时间
	return time.Duration(1.0 / lb.leakRate * float64(time.Second))
}

var (
	globalLimiter     Limiter
	ipLimiters        = make(map[string]Limiter)
	userLimiters      = make(map[string]Limiter)
	limitersMutex     sync.RWMutex
	cleanupTicker     *time.Ticker
)

// Init 初始化限流器
func Init(cfg *config.Config) {
	if !cfg.RateLimit.Enabled {
		return
	}

	// 全局限流器
	if cfg.RateLimit.Global.Enabled {
		switch cfg.RateLimit.Strategy {
		case "token_bucket":
			globalLimiter = NewTokenBucket(cfg.RateLimit.Global.Burst, cfg.RateLimit.Global.QPS)
		case "leaky_bucket":
			globalLimiter = NewLeakyBucket(cfg.RateLimit.Global.Burst, cfg.RateLimit.Global.QPS)
		}
	}

	// 定期清理不活跃的限流器
	cleanupTicker = time.NewTicker(5 * time.Minute)
	go func() {
		for range cleanupTicker.C {
			cleanupLimiters()
		}
	}()
}

// getOrCreateIPLimiter 获取或创建 IP 限流器
func getOrCreateIPLimiter(ip string, cfg config.RateLimitRule) Limiter {
	limitersMutex.RLock()
	limiter, exists := ipLimiters[ip]
	limitersMutex.RUnlock()

	if !exists {
		limitersMutex.Lock()
		defer limitersMutex.Unlock()

		// 双重检查
		limiter, exists = ipLimiters[ip]
		if !exists {
			var strategy string
			if cfg.QPS > 100 {
				strategy = "token_bucket"
			} else {
				strategy = "leaky_bucket"
			}

			switch strategy {
			case "token_bucket":
				limiter = NewTokenBucket(cfg.Burst, cfg.QPS)
			case "leaky_bucket":
				limiter = NewLeakyBucket(cfg.Burst, cfg.QPS)
			}

			ipLimiters[ip] = limiter
		}
	}

	return limiter
}

// getOrCreateUserLimiter 获取或创建用户限流器
func getOrCreateUserLimiter(userID string, cfg config.RateLimitRule) Limiter {
	limitersMutex.RLock()
	limiter, exists := userLimiters[userID]
	limitersMutex.RUnlock()

	if !exists {
		limitersMutex.Lock()
		defer limitersMutex.Unlock()

		limiter, exists = userLimiters[userID]
		if !exists {
			var strategy string
			if cfg.QPS > 100 {
				strategy = "token_bucket"
			} else {
				strategy = "leaky_bucket"
			}

			switch strategy {
			case "token_bucket":
				limiter = NewTokenBucket(cfg.Burst, cfg.QPS)
			case "leaky_bucket":
				limiter = NewLeakyBucket(cfg.Burst, cfg.QPS)
			}

			userLimiters[userID] = limiter
		}
	}

	return limiter
}

// cleanupLimiters 清理不活跃的限流器
func cleanupLimiters() {
	// TODO: 实现清理逻辑，移除长时间未使用的限流器
	// 简化版本，实际应该跟踪最后使用时间
}

// RateLimitMiddleware 限流中间件
func RateLimitMiddleware(rule config.RateLimitRule, cfg *config.Config) gin.HandlerFunc {
	return func(c *gin.Context) {
		// 全局限流
		if globalLimiter != nil && !globalLimiter.Allow() {
			logger.Warnf("全局限流触发: %s", c.ClientIP())
			c.JSON(429, gin.H{
				"error": "请求过于频繁，请稍后再试",
			})
			c.Abort()
			return
		}

		// IP 限流
		if cfg.RateLimit.PerIP.Enabled {
			ipLimiter := getOrCreateIPLimiter(c.ClientIP(), cfg.RateLimit.PerIP)
			if !ipLimiter.Allow() {
				logger.Warnf("IP 限流触发: %s", c.ClientIP())
				c.JSON(429, gin.H{
					"error": "请求过于频繁，请稍后再试",
				})
				c.Abort()
				return
			}
		}

		// 用户限流
		if cfg.RateLimit.PerUser.Enabled {
			userID, exists := c.Get("user_id")
			if exists {
				userLimiter := getOrCreateUserLimiter(userID.(string), cfg.RateLimit.PerUser)
				if !userLimiter.Allow() {
					logger.Warnf("用户限流触发: %s", userID)
					c.JSON(429, gin.H{
						"error": "请求过于频繁，请稍后再试",
					})
					c.Abort()
					return
				}
			}
		}

		c.Next()
	}
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}


